#include "libavutil/avutil.h"
#include "libavutil/imgutils.h"
#include "libavformat/avformat.h"
#include "libavcodec/avcodec.h"
#include <stdio.h>
void encode(AVFormatContext *out_format_context, AVCodecContext *enc_ctx, AVFrame *frame, AVPacket *newpkt);

int main(int argc, char *argv[])
{
    int ret;
    char *input = NULL;
    char *output = NULL;
    AVOutputFormat *output_fmt;
    AVFormatContext *ofmt_ctx;
    AVCodec *pCodec;
    int width = 960;
    int height = 540;
    av_log_set_level(AV_LOG_INFO);
    if (argc < 3)
    {
        av_log(NULL, AV_LOG_ERROR, "缺少输入和输出文件");
        return -1;
    }
    input = argv[1];
    output = argv[2];

    //第一步获取AVOutputFormat
    output_fmt = av_guess_format(NULL, output, NULL);

    //第二步创建AVFormatContext
    ofmt_ctx = avformat_alloc_context();

    if (avformat_alloc_output_context2(&ofmt_ctx, output_fmt, output_fmt->name, output) < 0)
    {
        av_log(NULL, AV_LOG_ERROR, "fail to alloc output context\n");
        ret = -1;
        goto __ERROR;
    }


    //第三步创建编码器
    pCodec =  avcodec_find_encoder_by_name("libx264");
    AVCodecContext *pCodecCtx = avcodec_alloc_context3(pCodec);
    pCodecCtx->codec_type = AVMEDIA_TYPE_VIDEO;
    pCodecCtx->profile = FF_PROFILE_H264_MAIN;
    pCodecCtx->level = 50; //标识level 是5.0
    pCodecCtx->width = width;
    pCodecCtx->height = height;
    pCodecCtx->gop_size = 250;
    pCodecCtx->keyint_min = 25;
    pCodecCtx->max_b_frames = 0;
    pCodecCtx->has_b_frames = 2;
    pCodecCtx->refs = 1;
    pCodecCtx->pix_fmt = AV_PIX_FMT_YUV420P;
    pCodecCtx->bit_rate = 600000;
    pCodecCtx->time_base = (AVRational){1, 25};
    pCodecCtx->framerate = (AVRational){25, 1};

    //第四步创建输出流并设置输出流的编码参数
    AVStream *out_stream = avformat_new_stream(ofmt_ctx, pCodec);
    avcodec_parameters_from_context(out_stream->codecpar, pCodecCtx);
    
    if (out_stream == NULL)
    {
        av_log(NULL, AV_LOG_ERROR, "fail to create new stream\n");
        ret = -1;
        goto __ERROR;
    }
    if (pCodec == NULL)
    {
        av_log(NULL, AV_LOG_ERROR, "fail to find codec\n");
        ret = -1;
        goto __ERROR;
    }

    // 第五步 打开编码器
    if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0)
    {
        av_log(NULL, AV_LOG_ERROR, "fail to open codec\n");
        ret = -1;
        goto __ERROR;
    }
    // avformat_find_stream_info(ofmt_ctx, NULL);
    av_dump_format(ofmt_ctx, 0, output, 1);

    // 第六步 获取输出h264文件的AVIOContext
    if (avio_open(&ofmt_ctx->pb, output, AVIO_FLAG_WRITE) < 0)
    {
        av_log(NULL, AV_LOG_ERROR, "fail to open output\n");
        ret = -1;
        goto __ERROR;
    }

    // 第七步写h264文件头
    if (avformat_write_header(ofmt_ctx, NULL) < 0)
    {
        av_log(NULL, AV_LOG_ERROR, "fail to write header");
        ret = -1;
        goto __ERROR;
    }

    FILE *fp = fopen(input, "rb");
    if (fp == NULL)
    {
        printf("fail to open file\n");
        ret = -1;
        goto __ERROR;
    }
    int base = 0;

    //第八步创建AVFrame 
    AVFrame *pframe = av_frame_alloc();
    pframe->format = pCodecCtx->pix_fmt;
    pframe->width = pCodecCtx->width;
    pframe->height = pCodecCtx->height;
    ret = av_frame_get_buffer(pframe, 32);
    if (ret < 0)
    {
        printf("av_frame_get_buffer fail error =  %s", av_err2str(ret));
        goto __ERROR;
    }
    //第九步 创建buffer用于存储文件读取出来的数据
    size_t buffer_size = av_image_get_buffer_size(pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height, 32);
    uint8_t *buffer = malloc(buffer_size);


    // 初始化输出AVPacket
    AVPacket newPacket;
    av_init_packet(&newPacket);

    //第十步读取文件内容，给编码器进行编码
    while (1)
    {
        //yuv格式yuv444 占用的空间是 width * height * 3,  yuv420 = width * height * 3 /2;
        int ret = fread(buffer, buffer_size, 1, fp);
        if (ret == 0)
        {
            break;
        }
        av_image_fill_arrays(pframe->data, pframe->linesize, buffer, pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height, 32);
        pframe->pts = base++;
        encode(ofmt_ctx, pCodecCtx, pframe, &newPacket);
    }

    // 第十一步 flush_encoder
    encode(ofmt_ctx, NULL, pframe, &newPacket);

    // 第十二步  av_write_trailer
    av_write_trailer(ofmt_ctx);

    //第十三步 释放资源
__ERROR:
    if (ofmt_ctx->pb)
    {
        avio_close(ofmt_ctx->pb);
    }

    if (pCodecCtx)
    {
        avcodec_close(pCodecCtx);
        avcodec_free_context(&pCodecCtx);
    }

    if (ofmt_ctx)
    {
        avformat_close_input(&ofmt_ctx);
        avformat_free_context(ofmt_ctx);
    }

    if (fp)
    {
        fclose(fp);
    }

    if (pframe)
    {
        av_frame_free(&pframe);
    }

    return ret;
}

void encode(AVFormatContext *out_format_context, AVCodecContext *enc_ctx, AVFrame *frame, AVPacket *newpkt)
{
    int ret = 0;
    if (!enc_ctx)
    {
        printf("encode  enc_ctx =  null");
        exit(1);
    }

    if (!newpkt)
    {
        printf("encode  newpkt =  null");
        exit(1);
    }

    if (!out_format_context)
    {
        printf("encode  newpkt =  null");
        exit(1);
    }

    if (frame)
    {
        printf("send frame to encoder, pts=%lld", frame->pts);
    }
    ret = avcodec_send_frame(enc_ctx, frame);

    if (ret < 0)
    {
        printf("encode  avcodec_send_frame  error");
        exit(1);
    }

    while (ret >= 0)
    {
        ret = avcodec_receive_packet(enc_ctx, newpkt);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
        {
            return;
        }
        else if (ret < 0)
        {
            printf("encode  avcodec_receive_packet error\n");
            exit(1);
        }
        av_write_frame(out_format_context, newpkt);
        av_packet_unref(newpkt);
    }
}
